"""
python fujian1_pmdarima_single_category.py
"""

import os
import json
import pandas as pd
from pmdarima import auto_arima
import logging
from pmdarima import ARIMA  # Import ARIMA directly

def predict_and_save_inventory(input_file_path, output_file_path, log_path=None):
    """
    根据输入的JSON文件中的库存数据，使用pmdarima的auto_arima进行预测，
    输出包含12个原始数据和3个预测数据的JSON文件。

    参数:
    - input_file_path: str, 输入JSON文件的路径
    - output_file_path: str, 输出JSON文件的路径
    - log_path: str, 日志文件的路径（可选，默认当前目录）
    """
    # 设置日志路径
    if log_path is None:
        log_path = os.getcwd()  # 默认使用当前目录

    log_file_path = os.path.join(log_path, 'prediction.log')
    logging.basicConfig(filename=log_file_path, level=logging.WARNING, format='%(asctime)s - %(levelname)s - %(message)s')

    try:
        # 读取输入 JSON 文件
        with open(input_file_path, 'r') as f:
            data = json.load(f)

        # 将数据转换为 DataFrame
        df = pd.DataFrame(data)
        df['date'] = pd.to_datetime(df['date'])
        df.set_index('date', inplace=True)

        # 确保有 12 个数据点进行训练
        if len(df) != 12:
            logging.warning("输入数据点不足 12 个，无法进行 ARIMA 预测。")
            return

        # 填充缺失值（如果有）
        df['inventory'] = df['inventory'].ffill()

        # 检查是否存在 NaN 值
        if df['inventory'].isnull().any():
            logging.error("数据包含 NaN 值，无法进行预测。")
            return

        # # 使用 auto_arima 选择模型
        # model = auto_arima(
        #     df['inventory'],
        #     seasonal=False,
        #     stepwise=True,
        #     suppress_warnings=True,
        # )

        
        model = ARIMA(order=(1, 1, 1))  # Initialize without y
        model.fit(df['inventory'])  # Fit the model with your data



        print("Selected ARIMA model parameters (p, d, q):", model.order)


        # 预测未来 3 个数据点
        forecast, conf_int = model.predict(n_periods=3, return_conf_int=True)
        
        print("预测出来的三个月的结果： ", forecast)

        # 生成未来日期
        forecast_dates = pd.date_range(start=df.index[-1] + pd.DateOffset(months=1), periods=3, freq='MS')

        # 创建预测结果 DataFrame
        forecast_df = pd.DataFrame({'date': forecast_dates, 'inventory': forecast})

        # 合并原始数据和预测数据
        result_df = pd.concat([df.reset_index(), forecast_df], ignore_index=True)

        # 将日期转换为字符串格式，以便 JSON 序列化
        result_df['date'] = result_df['date'].dt.strftime('%Y-%m-%dT00:00:00.000')

        # 转换为 JSON 格式并保存到输出文件
        result_json = result_df.to_dict(orient='records')
        with open(output_file_path, 'w') as f:
            json.dump(result_json, f, indent=4)

        print(f"预测结果已保存至 {output_file_path}")

    except Exception as e:
        logging.error(f"处理文件时发生错误：{e}")

# 示例用法
if __name__ == "__main__":
    input_file = r"..\fujian\fujian1\spline_then_smooth\smoothAfterSpline_category32.json"
    output_dir = r"..\fujian\fujian1\pmdarima_single_batch\unit_test_output"
    os.makedirs(output_dir, exist_ok=True)
    output_file = os.path.join(output_dir, "unit_test_output_category32.json")
    
    predict_and_save_inventory(input_file, output_file)
